from statsmodels.tsa.seasonal import seasonal_decompose
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import sys

def read_data(filename, start_hour, end_hour):
    data = pd.read_csv(filename)


    # 转换日期时间格式
    data['timestamp'] = pd.to_datetime(data['timestamp'], format='%Y/%m/%d %H:%M')
    date_time = data['timestamp']

    # 设定起始时间和结束时间
    start_time = pd.to_datetime(date_time.apply(lambda x: x.replace(hour=start_hour, minute=0)))
    end_time = pd.to_datetime(date_time.apply(lambda x: x.replace(hour=end_hour, minute=0)))

    # 选择指定时间段内的数据
    selected_data = data[(date_time >= start_time) & (date_time <= end_time)]

    return selected_data

if __name__ == "__main__":

    # 获取数据
    files = sys.argv
    filename = files[1]
    print('fileName:', filename)

    start_hour = int(files[2])
    end_hour = int(files[3])

    data = read_data(filename, start_hour, end_hour)
    total_rows = len(data)
    print('total rows:', total_rows)
    timestamp = np.array(data.timestamp)

    dfi = pd.DataFrame({
        'timestamp': timestamp,
        'Active_Power': data['Active_Power']
    })
    output_filename = f'.\\py\\result\\data{start_hour}-{end_hour}.csv'
#     output_filename = f'.\\data{start_hour}-{end_hour}.csv'
    print('output_filename:', output_filename)
    dfi.to_csv(output_filename, index=False)

    dfdata2020final = pd.DataFrame({
        'Active_Power': data['Active_Power']
    })
    data2020final_filename = '.\\py\\result\\data202122final.csv'
#     data2020final_filename = '.\\result\\data202122final.csv'

    dfdata2020final.to_csv(data2020final_filename, index=False)

    # 一层分解
    data_power = data['Active_Power'].tolist()

    # 进行季节趋势分解 （STL）
    data_decomp = seasonal_decompose(data_power, period=end_hour-start_hour+1)
    seasonal = data_decomp.seasonal
    trend = np.nan_to_num(data_decomp.trend)
    resid = np.nan_to_num(data_decomp.resid)

    def calculate_F(R, TorS):

        # Calculate variances
        var_R = np.var(R)
        var_T_plus_R = np.var(TorS + R)

        # Calculate F_T using the given formula
        F = max(0, 1 - (var_R / var_T_plus_R))

        return F

    F_T = calculate_F(resid, trend)
    F_S = calculate_F(resid, seasonal)
    print(F_T, F_S)

    df = pd.DataFrame({
        'timestamp': timestamp,
        'seasonal': seasonal,
        'trend': trend,
        'residual': resid,
    })
    output_filename = '.\\matlab\\forecast\\data_decomposednew.csv'
#     output_filename = '..\\matlab\\forecast\\data_decomposednew.csv'

    new_row = [total_rows, start_hour, end_hour, 0.8]
    new_row_df = pd.DataFrame([new_row], columns=df.columns)
    df = pd.concat([df, new_row_df], ignore_index=True)

    df.to_csv(output_filename, index=False)


    # 将数据存入 DataFrame
    # df = pd.DataFrame({
    #     'timestamp': timestamp,
    #     'seasonal': seasonal,
    #     'trend': trend,
    #     'residual': resid,
    #     'Temperature': data['Temperature'],
    #     'Humidity': data['Humidity'],
    #     'Horizontal_Radiation': data['Horizontal_Radiation'],
    #     'Scattered_Radiation': data['Scattered_Radiation'],
    #     'Precipitation': data['Precipitation'],
    # })
    #
    # # 将 DataFrame 导出为 CSV 文件
    # output_filename = '1.csv'
    # df.to_csv(output_filename, index=False)

